#ifndef FACTORY_H  
#define FACTORY_H


#include "trthelper.h"
#include "trtModelDeploy.h"
#include "common.h"
#include "yolov5_deploy.h"
#include "yolov8_deploy.h"
#include "yolov5-seg_deploy.h"
#include "yolov8-seg_deploy.h"
#include "yolov8-pose_deploy.h"
#include "yolov6-face_deploy.h"
class ModelFactory{
public:
    static std::unique_ptr<trtModelDeploy> createModel(trtModelType modelType, 
                                                    const std::string& modelPath, 
                                                    const std::string& classPath)
    {
        switch (modelType)
        {
        case YOLOV5:
            return std::make_unique<yolov5_deploy>(modelPath, classPath);
        case YOLOV8:
        case YOLOV11:
            return std::make_unique<yolov8_deploy>(modelPath, classPath);
        case YOLOV5_SEG:
            return std::make_unique<yolov5_seg_deploy>(modelPath, classPath);
        case YOLOV8_SEG:
        case YOLOV11_SEG:
            return std::make_unique<yolov8_seg_deploy>(modelPath, classPath);
        case YOLOV8_POSE:
        case YOLOV11_POSE:
            return std::make_unique<yolov8_pose_deploy>(modelPath, classPath);
        case YOLOV6_FACE:
            return std::make_unique<yolov6_face_deploy>(modelPath, classPath);
        default:
            throw std::runtime_error("Invalid model type");
        }
    } 
};
#endif